// Copyright (c) 2022 Huawei Device Co., Ltd. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "libcef/browser/storage/web_storage_impl.h"

#include <sstream>
#include "libcef/common/time_util.h"

#include "base/bind.h"
#include "base/logging.h"
#include "content/public/browser/browser_context.h"
#include "content/public/browser/storage_partition.h"
#include "third_party/blink/public/mojom/quota/quota_types.mojom.h"
#include "url/gurl.h"

class GetOriginsTask : public base::RefCountedThreadSafe<GetOriginsTask> {
 public:
  GetOriginsTask(CefWebStorageImpl::GetOriginsCallback callback,
                    storage::QuotaManager* quota_manager);
  void Run();

 private:
  friend class base::RefCountedThreadSafe<GetOriginsTask>;
  ~GetOriginsTask();
  void OnOriginsObtained(const std::set<url::Origin>& origins,
                         blink::mojom::StorageType type);
  void OnUsageAndQuotaObtained(const url::Origin& origin,
                               blink::mojom::QuotaStatusCode status_code,
                               int64_t usage,
                               int64_t quota);
  void CheckDone();
  void DoneOnUIThread();

  CefWebStorageImpl::GetOriginsCallback ui_callback_;
  scoped_refptr<storage::QuotaManager> quota_manager_;
  std::vector<std::string> origin_;
  std::vector<int64_t> usage_;
  std::vector<int64_t> quota_;
  size_t num_callbacks_to_wait_;
  size_t num_callbacks_received_;

  DISALLOW_COPY_AND_ASSIGN(GetOriginsTask);
};

GetOriginsTask::GetOriginsTask(CefWebStorageImpl::GetOriginsCallback callback,
    storage::QuotaManager* quota_manager) : ui_callback_(std::move(callback)), quota_manager_(quota_manager) {
  DCHECK(CEF_CURRENTLY_ON_UIT());
}

GetOriginsTask::~GetOriginsTask() {}

void GetOriginsTask::Run() {
  DCHECK(CEF_CURRENTLY_ON_UIT());
  CEF_POST_TASK(CEF_IOT,
                base::BindOnce(&storage::QuotaManager::GetOriginsModifiedBetween, quota_manager_,
                               blink::mojom::StorageType::kTemporary,
                               base::Time() /* Since beginning of time. */,
                               base::Time::Max() /* Until the end of times */,
                               base::BindOnce(&GetOriginsTask::OnOriginsObtained, this)));
}

void GetOriginsTask::OnOriginsObtained(
    const std::set<url::Origin>& origins,
    blink::mojom::StorageType type) {
  DCHECK(CEF_CURRENTLY_ON_IOT());
  num_callbacks_to_wait_ = origins.size();
  num_callbacks_received_ = 0u;
  for (const url::Origin& origin : origins) {
    quota_manager_->GetUsageAndQuota(origin, type,
                                     base::BindOnce(&GetOriginsTask::OnUsageAndQuotaObtained, this, origin));
  }
  CheckDone();
}

void GetOriginsTask::OnUsageAndQuotaObtained(
    const url::Origin& origin,
    blink::mojom::QuotaStatusCode status_code,
    int64_t usage,
    int64_t quota) {
  DCHECK(CEF_CURRENTLY_ON_IOT());
  if (status_code == blink::mojom::QuotaStatusCode::kOk) {
    origin_.push_back(origin.GetURL().spec());
    usage_.push_back(usage);
    quota_.push_back(quota);
  }
  ++num_callbacks_received_;
  CheckDone();
}

void GetOriginsTask::CheckDone() {
  DCHECK(CEF_CURRENTLY_ON_IOT());
  if (num_callbacks_received_ == num_callbacks_to_wait_) {
    CEF_POST_TASK(CEF_UIT,
                  base::BindOnce(&GetOriginsTask::DoneOnUIThread, this));
  } else if (num_callbacks_received_ > num_callbacks_to_wait_) {
    NOTREACHED();
  }
}

// This method is to avoid copying the 3 vector arguments into a bound callback.
void GetOriginsTask::DoneOnUIThread() {
  DCHECK(CEF_CURRENTLY_ON_UIT());
  if (ui_callback_) {
    std::move(ui_callback_).Run(origin_, usage_, quota_);
  }
}

CefBrowserContext* GetBrowserContext(const CefBrowserContext::Getter& getter) {
  DCHECK(!getter.is_null());
  // Will return nullptr if the BrowserContext has been destroyed.
  return getter.Run();
}

content::StoragePartition* GetStoragePartitionDelegate(CefBrowserContext* browser_context) {
  CEF_REQUIRE_UIT();
  return content::BrowserContext::GetDefaultStoragePartition(browser_context->AsBrowserContext());
}

storage::QuotaManager* GetQuotaManager(CefBrowserContext* browser_context) {
  CEF_REQUIRE_UIT();
  return GetStoragePartitionDelegate(browser_context)->GetQuotaManager();
}

void CefWebStorageImpl::DeleteAllData() {
  if (!ValidContext()) {
    StoreOrTriggerUIInitCallback(base::BindOnce(base::IgnoreResult(&CefWebStorageImpl::DeleteAllDataInternal), this));
    return;
  }
  return DeleteAllDataInternal();
}

void CefWebStorageImpl::DeleteAllDataInternal() {
  DCHECK(ValidContext());
  auto browser_context = GetBrowserContext(browser_context_getter_);
  if (!browser_context) {
    return;
  }
  GetStoragePartitionDelegate(browser_context)->ClearData(
      // Clear all web storage data except cookies.
      content::StoragePartition::REMOVE_DATA_MASK_APPCACHE |
          content::StoragePartition::REMOVE_DATA_MASK_FILE_SYSTEMS |
          content::StoragePartition::REMOVE_DATA_MASK_INDEXEDDB |
          content::StoragePartition::REMOVE_DATA_MASK_LOCAL_STORAGE |
          content::StoragePartition::REMOVE_DATA_MASK_WEBSQL,
      content::StoragePartition::QUOTA_MANAGED_STORAGE_MASK_TEMPORARY, GURL(),
      base::Time(), base::Time::Max(), base::DoNothing());
}

void CefWebStorageImpl::DeleteOrigin(const CefString& origin) {
  GURL gurl = GURL(origin.ToString16());
  if (!gurl.is_empty() && !gurl.is_valid())
    return;
  if (!ValidContext()) {
    StoreOrTriggerUIInitCallback(
        base::BindOnce(base::IgnoreResult(&CefWebStorageImpl::DeleteOriginInternal), this, gurl));
    return;
  }
  return DeleteOriginInternal(gurl);
}

void CefWebStorageImpl::DeleteOriginInternal(const GURL& origin) {
  DCHECK(ValidContext());
  DCHECK(origin.is_empty() || origin.is_valid());
  auto browser_context = GetBrowserContext(browser_context_getter_);
  if (!browser_context) {
    LOG(ERROR) << "CefWebStorageImpl::DeleteOriginInternal can not get browser_context.";
    return;
  }
  GetStoragePartitionDelegate(browser_context)->ClearDataForOrigin(
      // All (temporary) QuotaClient types.
      content::StoragePartition::REMOVE_DATA_MASK_APPCACHE |
          content::StoragePartition::REMOVE_DATA_MASK_FILE_SYSTEMS |
          content::StoragePartition::REMOVE_DATA_MASK_INDEXEDDB |
          content::StoragePartition::REMOVE_DATA_MASK_WEBSQL,
      content::StoragePartition::QUOTA_MANAGED_STORAGE_MASK_TEMPORARY, origin);
};

void CefWebStorageImpl::GetOrigins(CefRefPtr<CefGetOriginsCallback> callback) {
  if (!ValidContext()) {
    StoreOrTriggerUIInitCallback(base::BindOnce(
        base::IgnoreResult(&CefWebStorageImpl::GetOriginsInternal), this, callback));
    return;
  }
  return GetOriginsInternal(callback);
}

void CefWebStorageImpl::GetOriginsInternal(CefRefPtr<CefGetOriginsCallback> callback) {
  DCHECK(ValidContext());
  auto browser_context = GetBrowserContext(browser_context_getter_);
  if (!browser_context) {
    LOG(ERROR) << "CefWebStorageImpl::GetOriginsInternal can not get browser_context.";
    return;
  }
  GetOriginsCallback ui_callback =
      base::BindOnce(&CefWebStorageImpl::GetOriginsCallbackImpl,
                     weak_factory_.GetWeakPtr(), callback);
  base::MakeRefCounted<GetOriginsTask>(std::move(ui_callback),
                                       GetQuotaManager(browser_context))->Run();
}

void CefWebStorageImpl::GetOriginQuota(const CefString& origin,
                                       CefRefPtr<CefGetOriginUsageOrQuotaCallback> callback)
{
  if (!ValidContext()) {
    StoreOrTriggerUIInitCallback(base::BindOnce(
        base::IgnoreResult(&CefWebStorageImpl::GetOriginUsageAndQuotaInternal), this, origin, true, callback));
    return;
  }
  return GetOriginUsageAndQuotaInternal(origin, true, callback);
}

void CefWebStorageImpl::GetOriginUsage(const CefString& origin,
                                       CefRefPtr<CefGetOriginUsageOrQuotaCallback> callback)
{
  if (!ValidContext()) {
    StoreOrTriggerUIInitCallback(base::BindOnce(
        base::IgnoreResult(&CefWebStorageImpl::GetOriginUsageAndQuotaInternal), this, origin, false, callback));
    return;
  }
  return GetOriginUsageAndQuotaInternal(origin, false, callback);
}

void CefWebStorageImpl::OnUsageAndQuotaObtained(
    CefWebStorageImpl::QuotaUsageCallback ui_callback,
    blink::mojom::QuotaStatusCode status_code,
    int64_t usage,
    int64_t quota) {
  DCHECK(CEF_CURRENTLY_ON_IOT());
  if (status_code != blink::mojom::QuotaStatusCode::kOk) {
    usage = 0;
    quota = 0;
  }
  CEF_POST_TASK(CEF_UIT,
                base::BindOnce(std::move(ui_callback), usage, quota));
}

void CefWebStorageImpl::GetOriginUsageAndQuotaInternal(
    const CefString& origin,
    bool is_quota,
    CefRefPtr<CefGetOriginUsageOrQuotaCallback> callback)
{
  DCHECK(CEF_CURRENTLY_ON_UIT());
  auto browser_context = GetBrowserContext(browser_context_getter_);
  if (!browser_context) {
    LOG(ERROR) << "CefWebStorageImpl::GetOriginsInternal can not get browser_context.";
    return;
  }
  CefWebStorageImpl::QuotaUsageCallback ui_callback =
      base::BindOnce(&CefWebStorageImpl::CefGetOriginUsageOrQuotaCallbackImpl,
                     weak_factory_.GetWeakPtr(), callback, is_quota);
  CEF_POST_TASK(CEF_IOT,
                base::BindOnce(&storage::QuotaManager::GetUsageAndQuota, GetQuotaManager(browser_context), 
                url::Origin::Create(GURL(origin.ToString16())),
                blink::mojom::StorageType::kTemporary,
                base::BindOnce(&CefWebStorageImpl::OnUsageAndQuotaObtained, this, std::move(ui_callback))));
}

void RunAsyncCompletionOnUIThread(CefRefPtr<CefCompletionCallback> callback) {
  if (!callback.get())
    return;
  CEF_POST_TASK(CEF_UIT,
                base::Bind(&CefCompletionCallback::OnComplete, callback.get()));
}

void CefWebStorageImpl::Initialize(
    CefBrowserContext::Getter browser_context_getter,
    CefRefPtr<CefCompletionCallback> callback) {
  CEF_REQUIRE_UIT();
  DCHECK(!initialized_);
  DCHECK(!browser_context_getter.is_null());
  DCHECK(browser_context_getter_.is_null());
  browser_context_getter_ = browser_context_getter;
  initialized_ = true;
  if (!init_callbacks_.empty()) {
    for (auto& callback : init_callbacks_) {
      std::move(callback).Run();
    }
    init_callbacks_.clear();
  }
  RunAsyncCompletionOnUIThread(callback);
}

void CefWebStorageImpl::StoreOrTriggerUIInitCallback(
    base::OnceClosure callback) {
  if (!CEF_CURRENTLY_ON_UIT()) {
    CEF_POST_TASK(CEF_UIT,
                  base::BindOnce(&CefWebStorageImpl::StoreOrTriggerUIInitCallback, this, std::move(callback)));
    return;
  }
  if (initialized_) {
    std::move(callback).Run();
  } else {
    init_callbacks_.emplace_back(std::move(callback));
  }
}

bool CefWebStorageImpl::ValidContext() const {
  return CEF_CURRENTLY_ON_UIT() && initialized_;
}

void CefWebStorageImpl::GetOriginsCallbackImpl(
    CefRefPtr<CefGetOriginsCallback> callback,
    const std::vector<std::string>& origin,
    const std::vector<int64_t>& usage,
    const std::vector<int64_t>& quota) {
  if (!callback.get()) {
    return;
  }
  int32_t nums = origin.size();
  std::vector<CefString> origins;
  std::vector<CefString> usages;
  std::vector<CefString> quotas;
  std::string res;
  for (size_t i = 0; i < nums; i++) {
    origins.push_back(CefString(origin[i]));
    res = std::to_string(usage[i]);
    usages.push_back(CefString(res));
    res = std::to_string(quota[i]);
    quotas.push_back(CefString(res));
  }

  callback.get()->OnOrigins(origins);
  callback.get()->OnQuotas(quotas);
  callback.get()->OnUsages(usages);
  callback.get()->OnComplete();
}

void CefWebStorageImpl::CefGetOriginUsageOrQuotaCallbackImpl(
    CefRefPtr<CefGetOriginUsageOrQuotaCallback> callback,
    bool is_quota,
    int64_t usage,
    int64_t quota) {
  if (!callback.get()) {
    return;
  }
  int64_t result = is_quota == true ? quota : usage;
  callback.get()->OnComplete(result);
}

// static
CefRefPtr<CefWebStorage> CefWebStorage::GetGlobalManager(
    CefRefPtr<CefCompletionCallback> callback) {
  CefRefPtr<CefRequestContext> context = CefRequestContext::GetGlobalContext();
  return context ? context->GetWebStorage(callback) : nullptr;
}
